from torchvision.datasets import CIFAR100
from torchvision.transforms import Compose, ToTensor, Normalize

from data import Data


class Cifar100(Data):
    def __init__(self, n_clients, batch_size, alpha=1000, path='../', flag=True):
        super().__init__()
        transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
        self.dataset = CIFAR100(root=path, train=True, transform=transform)
        self.validate = CIFAR100(root=path, train=False, transform=transform)
        self.trainLoader, self.client_nums, self.total = \
            self.train_loader(alpha, n_clients, batch_size, flag)
        self.validationLoader = self.validate_loader(batch_size)

    def __str__(self):
        return "CIFAR100"
